# step 1: load documents from web pages
from langchain_community.document_loaders import WebBaseLoader

urls = [
    # "https://news.sina.com.cn/zx/gj/2025-07-16/doc-inffqtxh1722397.shtml",
    "https://k.sina.com.cn/article_1699432410_654b47da02001in8k.html",
    # "https://news.sina.com.cn/gov/2025-07-16/doc-inffqiin1921291.shtml",
]

docs = [WebBaseLoader(url).load() for url in urls]

# step 2: split documents into chunks
from langchain_text_splitters import RecursiveCharacterTextSplitter

docs_list = [item for sublist in docs for item in sublist]

text_splitter = RecursiveCharacterTextSplitter.from_tiktoken_encoder(
    chunk_size=100, chunk_overlap=50
)
doc_splits = text_splitter.split_documents(docs_list)

# step 3: create vector store and retriever
from langchain_core.vectorstores import InMemoryVectorStore
# from langchain_openai import OpenAIEmbeddings
# embed = OpenAIEmbeddings()

# faked embedding model for unit testing purposes.
# from langchain_core.embeddings import FakeEmbeddings
# embed = FakeEmbeddings(size=100)

# use ollama embedding model for unit testing purposes.
from langchain_ollama import OllamaEmbeddings
embed = OllamaEmbeddings(model="bge-m3:latest")

vectorstore = InMemoryVectorStore.from_documents(documents=doc_splits, embedding=embed)
retriever = vectorstore.as_retriever()

# step 4: create retriever tool and agent
from langchain.tools.retriever import create_retriever_tool

retriever_tool = create_retriever_tool(
    retriever,
    "retrieve_blog_posts",
    # "Search and return information about Lilian Weng blog posts.",
    "查找并返回关于链博会展的相关信息",
)

# test retriever tool
# print(retriever_tool.invoke({"query": "中央财政安排就业补助资金是多少"}))


# step 5: generate query or respond based on the current state

from langgraph.graph import MessagesState
# from langchain.chat_models import init_chat_model
# response_model = init_chat_model("openai:gpt-4.1", temperature=0)

from langchain_ollama import ChatOllama
response_model = ChatOllama(model="qwen2.5:7b-instruct-q5_K_S",temperature=0.2)

def generate_query_or_respond(state: MessagesState):
    """Call the model to generate a response based on the current state. Given
    the question, it will decide to retrieve using the retriever tool, or simply respond to the user.
    """
    response = (
        response_model
        .bind_tools([retriever_tool]).invoke(state["messages"])
    )
    return {"messages": [response]}

# test case 1: invoke generate_query_or_respond and response message without using the retriever tool.
# input = {"messages": [{"role": "user", "content": "hello!"}]}
# generate_query_or_respond(input)["messages"][-1].pretty_print()

# test case 2: invoke generate_query_or_respond and response message with using the retriever tool.
# input = {
#     "messages": [
#         {
#             "role": "user",
#             "content": "What does Lilian Weng say about types of reward hacking?",
#         }
#     ]
# }
# generate_query_or_respond(input)["messages"][-1].pretty_print()

# step 6: Add a conditional edge - grade documents using a binary score for relevance check
from pydantic import BaseModel, Field
from typing import Literal

GRADE_PROMPT = (
    "You are a grader assessing relevance of a retrieved document to a user question. \n "
    "Here is the retrieved document: \n\n {context} \n\n"
    "Here is the user question: {question} \n"
    "If the document contains keyword(s) or semantic meaning related to the user question, grade it as relevant. \n"
    "Give a binary score 'yes' or 'no' score to indicate whether the document is relevant to the question."
)

class GradeDocuments(BaseModel):
    """Grade documents using a binary score for relevance check."""
    binary_score: str = Field(
        description="Relevance score: 'yes' if relevant, or 'no' if not relevant"
    )

grader_model = ChatOllama(model="deepseek-r1:7b-qwen-distill-q4_K_M",temperature=0.2)

def grade_documents(
    state: MessagesState,
) -> Literal["generate_answer", "rewrite_question"]:
    """Determine whether the retrieved documents are relevant to the question."""
    question = state["messages"][0].content
    context = state["messages"][-1].content

    prompt = GRADE_PROMPT.format(question=question, context=context)
    response = (
        grader_model
        .with_structured_output(GradeDocuments).invoke(
            [{"role": "user", "content": prompt}]
        )
    )
    score = response.binary_score

    if score == "yes":
        return "generate_answer"
    else:
        return "rewrite_question"

# test case : invoke grade_documents and response message with using the retriever tool.
from langchain_core.messages import convert_to_messages
# input = {
#     "messages": convert_to_messages(
#         [
#             {
#                 "role": "user",
#                 "content": "What does Lilian Weng say about types of reward hacking?",
#             },
#             {
#                 "role": "assistant",
#                 "content": "",
#                 "tool_calls": [
#                     {
#                         "id": "1",
#                         "name": "retrieve_blog_posts",
#                         "args": {"query": "types of reward hacking"},
#                     }
#                 ],
#             },
#             {"role": "tool", "content": "meow", "tool_call_id": "1"},
#         ]
#     )
# }
# print(grade_documents(input))

# step 7: Build the rewrite_question node - rewrite question if the documents are not relevant
REWRITE_PROMPT = (
    "Look at the input and try to reason about the underlying semantic intent / meaning.\n"
    "Here is the initial question:"
    "\n ------- \n"
    "{question}"
    "\n ------- \n"
    "Only return the formulated and improved question"
)

def rewrite_question(state: MessagesState):
    """Rewrite the original user question."""
    messages = state["messages"]
    question = messages[0].content
    prompt = REWRITE_PROMPT.format(question=question)
    response = response_model.invoke([{"role": "user", "content": prompt}])
    return {"messages": [{"role": "user", "content": response.content}]}

# test case : invoke rewrite_question and response message with using the retriever tool.
# input = {
#     "messages": convert_to_messages(
#         [
#             {
#                 "role": "user",
#                 "content": "What does Lilian Weng say about types of reward hacking?",
#             },
#             {
#                 "role": "assistant",
#                 "content": "",
#                 "tool_calls": [
#                     {
#                         "id": "1",
#                         "name": "retrieve_blog_posts",
#                         "args": {"query": "types of reward hacking"},
#                     }
#                 ],
#             },
#             {"role": "tool", "content": "meow", "tool_call_id": "1"},
#         ]
#     )
# }

# response = rewrite_question(input)
# print(response["messages"][-1]["content"])

# step 8: Build the generate_answer node - generate an answer based on the retrieved documents
GENERATE_PROMPT = (
    "You are an assistant for question-answering tasks. "
    "Use the following pieces of retrieved context to answer the question. "
    "If you don't know the answer, just say that you don't know. "
    "Use three sentences maximum and keep the answer concise.\n"
    "Question: {question} \n"
    "Context: {context}"
)

def generate_answer(state: MessagesState):
    """Generate an answer."""
    question = state["messages"][0].content
    context = state["messages"][-1].content
    prompt = GENERATE_PROMPT.format(question=question, context=context)
    response = response_model.invoke([{"role": "user", "content": prompt}])
    return {"messages": [response]}

# test case : invoke generate_answer and response message with using the retriever tool.
# input = {
#     "messages": convert_to_messages(
#         [
#             {
#                 "role": "user",
#                 "content": "What does Lilian Weng say about types of reward hacking?",
#             },
#             {
#                 "role": "assistant",
#                 "content": "",
#                 "tool_calls": [
#                     {
#                         "id": "1",
#                         "name": "retrieve_blog_posts",
#                         "args": {"query": "types of reward hacking"},
#                     }
#                 ],
#             },
#             {
#                 "role": "tool",
#                 "content": "reward hacking can be categorized into two types: environment or goal misspecification, and reward tampering",
#                 "tool_call_id": "1",
#             },
#         ]
#     )
# }

# response = generate_answer(input)
# response["messages"][-1].pretty_print()

# step 9: Build the workflow graph - define the nodes and edges of the workflow graph.  Define the nodes we will cycle between and the edges taken after the `action` node is called.  Compile the graph.  Define the nodes and edges of the workflow graph.  Define the nodes we will cycle between and the edges taken after the `action` node is called.  Compile the graph.  Define the nodes and edges of the workflow graph.  Define the nodes we will cycle
from langgraph.graph import StateGraph, START, END
from langgraph.prebuilt import ToolNode
from langgraph.prebuilt import tools_condition

workflow = StateGraph(MessagesState)

# Define the nodes we will cycle between
workflow.add_node(generate_query_or_respond)
workflow.add_node("retrieve", ToolNode([retriever_tool]))
workflow.add_node(rewrite_question)
workflow.add_node(generate_answer)

# Start with a generate_query_or_respond and determine if we need to call retriever_tool
workflow.add_edge(START, "generate_query_or_respond")

# Decide whether to retrieve
workflow.add_conditional_edges(
    "generate_query_or_respond",
    # Assess LLM decision (call `retriever_tool` tool or respond to the user)
    tools_condition,
    {
        # Translate the condition outputs to nodes in our graph
        "tools": "retrieve",
        END: END,
    },
)

# Edges taken after the `action` node is called.
workflow.add_conditional_edges(
    "retrieve",
    # Assess agent decision
    grade_documents,
)
# If relevant, proceed to generate_answer and generate final response using the ToolMessage with the retrieved document context
workflow.add_edge("generate_answer", END)
# If not relevant, rewrite the question using rewrite_question and then call generate_query_or_respond again
workflow.add_edge("rewrite_question", "generate_query_or_respond")

# Compile
graph = workflow.compile()

# step 10: Run the workflow graph - run the workflow graph and print the output.
for chunk in graph.stream(
    {
        "messages": [
            {
                "role": "user",
                "content": "链博会都有多少家企业参加了?",
            }
        ]
    }
):
    for node, update in chunk.items():
        print("Update from node", node)
        update["messages"][-1].pretty_print()
        print("\n\n")
